{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9981f4b0",
   "metadata": {},
   "source": [
    "# Tutorial Notebook 10: Finetuning for Perturbation Response Prediction\n",
    "\n",
    "In this tutorial, we will demonstrate how to finetune a Cell2Sentence (C2S) model for perturbation response prediction tasks. This is a critical task in single-cell analysis, where the goal is to predict how a cell's gene expression profile changes in response to a specific perturbation (e.g., a genetic knockout or a drug treatment).\n",
    "\n",
    "We will treat this as a \"translation\" task in natural language: translating a cell (in cell sentence format) from its basal (control) state to its perturbed state, conditioned on the perturbation applied.\n",
    "\n",
    "At a high level, we will:\n",
    "1. Load a public single-cell perturbation dataset.\n",
    "2. Write a custom prompt template for perturbation prediction.\n",
    "3. Subclass the `PromptFormatter` class to create pairs of control and perturbed cells.\n",
    "4. Finetune a pretrained C2S-Scale model on this new task.\n",
    "5. Generate a prediction with our new finetuned model to see it in action."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "139d5b0b",
   "metadata": {},
   "source": [
    "First, let's import the necessary libraries. We'll need standard data science libraries, single-cell analysis tools, and modules from the `cell2sentence` and `transformers` packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f0bf13c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Python built-in libraries\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\" \n",
    "os.environ[\"WORLD_SIZE\"] = \"1\"\n",
    "\n",
    "import pickle\n",
    "import random\n",
    "from datetime import datetime\n",
    "from collections import Counter, defaultdict\n",
    "\n",
    "# Third-party libraries\n",
    "from datasets import Dataset\n",
    "import numpy as np\n",
    "import torch\n",
    "from transformers import TrainingArguments, AutoModelForCausalLM\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Single-cell libraries\n",
    "import anndata\n",
    "import scanpy as sc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cb07037d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell2Sentence imports\n",
    "import cell2sentence as cs\n",
    "from cell2sentence.prompt_formatter import get_cell_sentence_str, PromptFormatter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "81ce734c",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 1234\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f6047f8",
   "metadata": {},
   "source": [
    "# Load Perturbation Data\n",
    "\n",
    "For this tutorial, you will need a single-cell dataset that includes both control and perturbed cells. The data should be in an `AnnData` object. Crucially, your `.obs` dataframe must contain:\n",
    "- A column that distinguishes control cells from perturbed cells, e.g., `adata.obs['condition']`\n",
    "    - Values can be something like 'control' or 'non-targeting' for control cells, and 'perturbed' or the perturbation name for the perturbed cells\n",
    "\n",
    "For this example, we use a public genetic perturbation dataset of Jurkat cells (Nadig et al., 2025). To use your own dataset, replace `DATA_PATH` with the path to your preprocessed data file.\n",
    "- Paper: https://www.nature.com/articles/s41588-025-02169-3\n",
    "\n",
    "<font color='red'>Ensure your data is preprocessed and normalized (e.g., using log1p transformation) before proceeding.</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "52754a57",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 21412 × 18080\n",
       "    obs: 'batch_var', 'cell_type', 'target_gene', 'gene_id', 'mitopercent', 'UMI_count'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Replace this with the actual path to your dataset, if using a custom dataset\n",
    "DATA_PATH = \"/home/sr2464/scratch/C2S_API_Testing/Data/jurkat.h5ad\"\n",
    "adata = anndata.read_h5ad(DATA_PATH)\n",
    "adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "590a4a9f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>batch_var</th>\n",
       "      <th>cell_type</th>\n",
       "      <th>target_gene</th>\n",
       "      <th>gene_id</th>\n",
       "      <th>mitopercent</th>\n",
       "      <th>UMI_count</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>cell_barcode</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>AAACCCAAGCACCAGA-42</th>\n",
       "      <td>jurkat42</td>\n",
       "      <td>jurkat</td>\n",
       "      <td>EIF4B</td>\n",
       "      <td>ENSG00000063046</td>\n",
       "      <td>0.037697</td>\n",
       "      <td>4828.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>AAACCCAAGCAGATAT-4</th>\n",
       "      <td>jurkat4</td>\n",
       "      <td>jurkat</td>\n",
       "      <td>DAXX</td>\n",
       "      <td>ENSG00000204209</td>\n",
       "      <td>0.063256</td>\n",
       "      <td>9675.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>AAACCCAAGCTAGATA-2</th>\n",
       "      <td>jurkat2</td>\n",
       "      <td>jurkat</td>\n",
       "      <td>METTL3</td>\n",
       "      <td>ENSG00000165819</td>\n",
       "      <td>0.069885</td>\n",
       "      <td>16985.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>AAACCCAAGCTGGAGT-37</th>\n",
       "      <td>jurkat37</td>\n",
       "      <td>jurkat</td>\n",
       "      <td>non-targeting</td>\n",
       "      <td>non-targeting</td>\n",
       "      <td>0.055775</td>\n",
       "      <td>20475.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>AAACCCAAGGGCCTCT-51</th>\n",
       "      <td>jurkat51</td>\n",
       "      <td>jurkat</td>\n",
       "      <td>non-targeting</td>\n",
       "      <td>non-targeting</td>\n",
       "      <td>0.047619</td>\n",
       "      <td>12642.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                    batch_var cell_type    target_gene          gene_id  \\\n",
       "cell_barcode                                                              \n",
       "AAACCCAAGCACCAGA-42  jurkat42    jurkat          EIF4B  ENSG00000063046   \n",
       "AAACCCAAGCAGATAT-4    jurkat4    jurkat           DAXX  ENSG00000204209   \n",
       "AAACCCAAGCTAGATA-2    jurkat2    jurkat         METTL3  ENSG00000165819   \n",
       "AAACCCAAGCTGGAGT-37  jurkat37    jurkat  non-targeting    non-targeting   \n",
       "AAACCCAAGGGCCTCT-51  jurkat51    jurkat  non-targeting    non-targeting   \n",
       "\n",
       "                     mitopercent  UMI_count  \n",
       "cell_barcode                                 \n",
       "AAACCCAAGCACCAGA-42     0.037697     4828.0  \n",
       "AAACCCAAGCAGATAT-4      0.063256     9675.0  \n",
       "AAACCCAAGCTAGATA-2      0.069885    16985.0  \n",
       "AAACCCAAGCTGGAGT-37     0.055775    20475.0  \n",
       "AAACCCAAGGGCCTCT-51     0.047619    12642.0  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.obs.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e07fcedb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "69\n"
     ]
    }
   ],
   "source": [
    "target_gene_counter = Counter(adata.obs['target_gene'])\n",
    "print(len(target_gene_counter))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ab4e3458",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('non-targeting', 12013),\n",
       " ('TFAM', 2555),\n",
       " ('EIF4B', 461),\n",
       " ('JAZF1', 265),\n",
       " ('HIRA', 233),\n",
       " ('SMARCB1', 204),\n",
       " ('TAF13', 200),\n",
       " ('KDM1A', 195),\n",
       " ('MBTPS1', 192),\n",
       " ('LZTR1', 192),\n",
       " ('METTL3', 177),\n",
       " ('DNMT1', 177),\n",
       " ('SDC1', 156),\n",
       " ('TADA1', 144),\n",
       " ('MED12', 143),\n",
       " ('USF2', 141),\n",
       " ('ARPC2', 140),\n",
       " ('PTPN1', 136),\n",
       " ('PHF10', 136),\n",
       " ('NDUFB4', 133)]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_gene_counter.most_common(20)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afec3e85",
   "metadata": {},
   "source": [
    "This data contains both control cells (non-targeting) as well as cells under different genetic knockouts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "38ebcf86",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>SAMD11</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>NOC2L</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KLHL17</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>PLEKHN1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>PERM1</th>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Empty DataFrame\n",
       "Columns: []\n",
       "Index: [SAMD11, NOC2L, KLHL17, PLEKHN1, PERM1]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.var.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cc318b3f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<21412x18080 sparse matrix of type '<class 'numpy.float32'>'\n",
       "\twith 69336411 stored elements in Compressed Sparse Row format>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6239a8bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.6483034, 1.6483034, 1.6483034, 1.6483034, 1.6483034, 1.6483034,\n",
       "       1.6483034, 1.6483034, 1.6483034, 1.6483034], dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# View a few nonzero values\n",
    "adata.X.data[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0e4b95d8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8.333305"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.X.max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "005cef4b",
   "metadata": {},
   "source": [
    "The expression is already preprocessed and log1p transformed. Typically, a log1p transform with base=10 would be used for Cell2Sentence training if we wish to convert generated cell sentences back to expression vectors, since log base=10 gives a better linear relationship between log rank and log expression in many single-cell datasets.\n",
    "\n",
    "For this tutorial, we will use this processed data as is, to train our model to generate target cell sentences."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1c99514",
   "metadata": {},
   "source": [
    "# Cell2Sentence Conversion\n",
    "\n",
    "Now, we will convert the gene expression data in our `AnnData` object into cell sentences. This process creates a Hugging Face Arrow dataset, which is used in our LLM training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e0ae0bdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['batch_var',\n",
       " 'cell_type',\n",
       " 'target_gene',\n",
       " 'gene_id',\n",
       " 'mitopercent',\n",
       " 'UMI_count']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# We'll keep all relevant columns for our new task\n",
    "adata_obs_cols_to_keep = adata.obs.columns.tolist()\n",
    "adata_obs_cols_to_keep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "91e60261",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 21412/21412 [00:11<00:00, 1899.14it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['cell_name', 'cell_sentence', 'batch_var', 'cell_type', 'target_gene', 'gene_id', 'mitopercent', 'UMI_count'],\n",
       "    num_rows: 21412\n",
       "})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create Arrow dataset and vocabulary\n",
    "arrow_ds, vocabulary = cs.CSData.adata_to_arrow(\n",
    "    adata=adata, \n",
    "    random_state=SEED, \n",
    "    sentence_delimiter=' ',\n",
    "    label_col_names=adata_obs_cols_to_keep\n",
    ")\n",
    "arrow_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3efc581c",
   "metadata": {},
   "source": [
    "# Custom Prompt Formatting for Perturbation Prediction\n",
    "\n",
    "This is the core of our tutorial. For this dataset, we have a single large pool of control cells (labeled 'non-targeting') and multiple groups of perturbed cells, each with a specific `target_gene`. Our goal is to create training pairs where each perturbed cell is matched with a randomly selected control cell. Note that for different perturbation applications, there may be better ways of pairing control and perturbed cells.\n",
    "\n",
    "We will define a custom prompt structure that frames our task for the LLM. The input will contain the **control cell sentence** and the **perturbation name**. The model's expected output (the response) will be the **perturbed cell sentence**.\n",
    "\n",
    "First, let's define our prompt templates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "adceba43",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The input provides the control cell and the perturbation, asking for the perturbed result.\n",
    "custom_input_prompt_template = \"\"\"Given the following cell sentence of {num_genes} expressed genes representing a cell's basal state, predict the cell sentence after applying the perturbation: {perturbation_name}.\n",
    "Control cell sentence: {control_cell_sentence}.\n",
    "\n",
    "Perturbed cell sentence:\"\"\"\n",
    "\n",
    "# The answer is simply the target cell sentence.\n",
    "answer_template = \"{perturbed_cell_sentence}.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8afe106b",
   "metadata": {},
   "source": [
    "To apply this template, we need to create pairs of (control cell, perturbed cell) for each perturbation. We'll create a custom `PerturbationPromptFormatter` by subclassing the base `PromptFormatter`.\n",
    "\n",
    "Our custom `format_hf_ds` function will:\n",
    "1.  First, iterate through the entire dataset to create a list of all control cell indices.\n",
    "2.  Simultaneously, it will group the indices of all perturbed cells into a dictionary, with the perturbation name (`target_gene`) as the key.\n",
    "3.  Then, it will iterate through each perturbation group and, for every perturbed cell, randomly select a control cell from the global pool to form a pair.\n",
    "4.  Finally, it will format these pairs into the input prompts and expected responses for the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "84fdc5ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "class PerturbationPromptFormatter(PromptFormatter):\n",
    "    def __init__(self,\n",
    "        task_name,\n",
    "        input_prompt,\n",
    "        answer_template,\n",
    "        top_k_genes, \n",
    "        perturbation_col='target_gene',\n",
    "        control_label='non-targeting'\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Initializes the custom prompt formatter.\n",
    "\n",
    "        Args:\n",
    "            task_name (str): The name for this task.\n",
    "            input_prompt (str): The template for the model's input.\n",
    "            answer_template (str): The template for the model's expected response.\n",
    "            top_k_genes (int): The number of top genes to include in the cell sentence.\n",
    "            perturbation_col (str): The column name in the dataset that contains perturbation info.\n",
    "            control_label (str): The label used to identify control cells in the perturbation_col.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.task_name = task_name\n",
    "        self.input_prompt = input_prompt\n",
    "        self.answer_template = answer_template\n",
    "        self.top_k_genes = top_k_genes\n",
    "        self.perturbation_col = perturbation_col\n",
    "        self.control_label = control_label\n",
    "        assert top_k_genes > 0, \"'top_k_genes' must be an integer > 0\"\n",
    "\n",
    "    def format_hf_ds(self, hf_ds):\n",
    "        \"\"\"\n",
    "        Custom formatting function for perturbation prediction. This function creates pairs of\n",
    "        (control, perturbed) cells by sampling from a global pool of control cells.\n",
    "        \"\"\"\n",
    "        model_inputs_list = []\n",
    "        responses_list = []\n",
    "        \n",
    "        # 1. Separate all cells into a global control pool and a dict of perturbed cells\n",
    "        control_indices = []\n",
    "        pert_to_indices = defaultdict(list)\n",
    "\n",
    "        print(\"Grouping cells by perturbation and identifying global controls...\")\n",
    "        for i, sample in enumerate(hf_ds):\n",
    "            if sample[self.perturbation_col] == self.control_label:\n",
    "                control_indices.append(i)\n",
    "            else:\n",
    "                pert_to_indices[sample[self.perturbation_col]].append(i)\n",
    "        \n",
    "        assert len(control_indices) > 0, \"No control cells found. Cannot create pairs.\"\n",
    "        print(f\"Found {len(control_indices)} control cells.\")\n",
    "        print(f\"Found {len(pert_to_indices)} unique perturbations.\")\n",
    "\n",
    "        # 2. Create prompt-response pairs by iterating through perturbed cells\n",
    "        print(\"Creating control-perturbed pairs...\")\n",
    "        for pert_name, perturbed_indices in tqdm(pert_to_indices.items()):\n",
    "            for perturbed_idx in perturbed_indices:\n",
    "                # Pair each perturbed cell with a random control cell from the global pool\n",
    "                control_idx = random.choice(control_indices)\n",
    "                \n",
    "                control_sample = hf_ds[control_idx]\n",
    "                perturbed_sample = hf_ds[perturbed_idx]\n",
    "\n",
    "                # Format control cell sentence\n",
    "                control_sentence, num_genes_str = get_cell_sentence_str(\n",
    "                    control_sample,\n",
    "                    num_genes=self.top_k_genes\n",
    "                )\n",
    "                # Format perturbed cell sentence\n",
    "                perturbed_sentence, _ = get_cell_sentence_str(\n",
    "                    perturbed_sample,\n",
    "                    num_genes=self.top_k_genes\n",
    "                )\n",
    "\n",
    "                # Format the model input string using the perturbation name\n",
    "                model_input_str = self.input_prompt.format(\n",
    "                    num_genes=num_genes_str,\n",
    "                    perturbation_name=pert_name,\n",
    "                    control_cell_sentence=control_sentence\n",
    "                )\n",
    "                \n",
    "                # Format the response string\n",
    "                response_str = self.answer_template.format(\n",
    "                    perturbed_cell_sentence=perturbed_sentence\n",
    "                )\n",
    "\n",
    "                model_inputs_list.append(model_input_str)\n",
    "                responses_list.append(response_str)\n",
    "\n",
    "        # Create the final Hugging Face Dataset\n",
    "        ds_split_dict = {\n",
    "            \"sample_type\": [self.task_name] * len(model_inputs_list),\n",
    "            \"model_input\": model_inputs_list,\n",
    "            \"response\": responses_list,\n",
    "        }\n",
    "        ds = Dataset.from_dict(ds_split_dict)\n",
    "        return ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e431c369",
   "metadata": {},
   "source": [
    "Let's instantiate our custom formatter and test it on a small sample of our data to see the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "68228040",
   "metadata": {},
   "outputs": [],
   "source": [
    "task_name = \"perturbation_prediction\"\n",
    "prompt_formatter = PerturbationPromptFormatter(\n",
    "    task_name=task_name,\n",
    "    input_prompt=custom_input_prompt_template,\n",
    "    answer_template=answer_template,\n",
    "    top_k_genes=200 # Using top 200 genes for this example. For real applications, ideal to use all nonzero expressed genes if possible.\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "51258c12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Grouping cells by perturbation and identifying global controls...\n",
      "Found 12013 control cells.\n",
      "Found 68 unique perturbations.\n",
      "Creating control-perturbed pairs...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 68/68 [00:03<00:00, 20.85it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['sample_type', 'model_input', 'response'],\n",
       "    num_rows: 9399\n",
       "})"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Format the dataset\n",
    "formatted_ds = prompt_formatter.format_hf_ds(arrow_ds)\n",
    "formatted_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95e9668c",
   "metadata": {},
   "source": [
    "Note that if you want to do a train/test split of the data, separating out a split of control cells and holdout perturbed cells / entire perturbations can be done before formatting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f8b0f343",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Formatted Sample ---\n",
      "#----Model input:----#\n",
      "Given the following cell sentence of 200 expressed genes representing a cell's basal state, predict the cell sentence after applying the perturbation: EIF4B.\n",
      "Control cell sentence: MT-ATP6 MT-CO3 TMSB4X MT-ND4 MT-ND1 MT-CO2 ACTB MT-ND2 MT-ND3 MT-CYB RACK1 HSP90AB1 HSP90AA1 NME2 HSPD1 H2AFZ YBX1 ACTG1 FTH1 TUBA1B TUBB LDHB PFN1 NCL FTL PRDX1 EEF1B2 HIST1H4C BTF3 ADA CFL1 CHCHD2 SET ENO1 UBA52 H3F3A ARHGDIB EIF1 EEF1G PPP1R14B GSTP1 CALR HINT1 SERF2 HSPE1 XRCC5 TMSB10 STMN1 CD3D B2M EEF2 MT-ND6 PPIA HIST1H1D MIF H3F3B ANP32B HNRNPU UBC SERBP1 CDK6 EIF4A1 TYMS UQCRH PSMA7 HNRNPA2B1 SFPQ GPX4 HMGB2 PRRC2C SIVA1 FUS SUB1 LCP1 SLC25A3 SNRPB SRSF3 CSDE1 PCLAF EIF3G EIF3A MYL6 HES4 HMGN1 SLC25A5 COX4I1 MKI67 NDUFA13 HNRNPD PEBP1 PPP1CA GABARAP PRDX5 HNRNPA3 TIMM13 PKM POMP PCNA NUDC STOML2 GADD45GIP1 TMA7 PPA1 CORO1A HNRNPK ARPP19 UQCRB ATP5F1B SNRPD1 DDX5 ARPC3 ISG15 PSMA4 SEC61B COX8A PSMA1 FABP5 TMEM160 SNRPE PRPF40A MT-ND5 ODC1 COX6A1 SH3BGRL3 HMGA1 NUCKS1 COX7A2 ATP5MC2 NOP56 HNRNPC KTN1 SRSF7 SOX4 CLIC1 PSMD2 NDUFS6 ALYREF SUMO2 GUK1 SSBP1 YWHAQ BRD2 SRSF9 ERH TOMM6 PSMC5 CCT2 CCT4 IGFBP2 C1QBP ATP5F1E ARL6IP4 C12ORF57 SNU13 METTL5 CCNI ATP5MC1 CCT8 DDX46 TMEM208 JPT1 SPCS1 FDFT1 PA2G4 HNRNPR HIST1H1B CENPF SNRPF CCT3 ARPC1B COX6C DBI RBM3 PRKDC EIF3I C12ORF75 DIAPH1 EIF5 MARCKSL1 IFI16 EIF4G2 MT-ND4L PHB CCT6A ARPC5 NUCB2 LIME1 POLR2F NME1 PPM1G VIM MANF UBL5 COPS6 PSMB6 SRRM1 SUPT16H SRM SELENOH AP2M1.\n",
      "\n",
      "Perturbed cell sentence: \n",
      "\n",
      "#----Response:----#\n",
      "MT-ATP6 MT-CO3 MT-CO2 TUBA1B HSP90B1 MT-CYB MT-ND4 ACTB TMSB4X HSP90AA1 HSP90AB1 CENPF MT-ND6 TRBC1 MT-ND2 HNRNPA2B1 PPIA H3F3A HSPD1 GSTP1 TUBB LDHB CHCHD2 SET MIF ADA RACK1 EEF1G EEF1B2 BTF3 CD3G NME2 NCL HSPA8 MT-ND1 XRCC5 EIF4A1 NUCKS1 CCNB1 PRDX1 EDF1 H2AFY RNASEH2C CDK6 FABP5 FTH1 HNRNPD COX4I1 KPNA2 CDC123 UBA52 B2M MKI67 SNRPB TPM4 ARPC5 HINT1 ITM2B CORO1A ATP5MC3 H2AFZ PMEPA1 MED28 ARHGDIB MYL12A YBX3 ASPM H2AFX ATP5PF RBM3 SLC25A5 HNRNPU ZNRF1 PFN1 MYO7B RNF213 ACTG1 YBX1 C4ORF48 HSPA5 ITM2A SRSF2 KHDRBS1 EIF2S3 DDX39A SFPQ TECR CCT2 HSPE1 SMC4 CANX DDX5 CSDE1 ATP5F1D TCF12 C1QBP CD3D FYB1 MT-ND3 TMBIM6 SERF2 NKTR CYC1 NDUFA4 DIAPH1 PALM2-AKAP2 FTL MT-ND5 SOX4 H3F3B LEF1 GTF3A GABARAP XRN2 LSM4 CLIC1 NIPBL CHORDC1 STMN1 SIX6 PGK1 BLOC1S6 UBE2I HMGB2 IRF1 CD3E SRRM1 NDUFA7 CFL1 KNL1 GUK1 NSA2 PRPF8 IFI16 SNRPE RNF220 EIF4G2 SETX STT3B RUFY1 XRCC6 NAE1 PNN NOP10 IK SRP19 SUB1 MSI2 TXNL4A PFDN5 SRM PCBP2 NCAPG DDX3X VPS26B TCOF1 JPT1 EIF1 MAD2L1 PSMD2 CUL4A MARCKSL1 NUCB2 NDUFA2 ATP5MG TUBA1A LBR CDK1 MAP4K4 COPS5 NOP58 SAMD1 ENO1 PFDN2 MTCH2 KMT5A TMEM160 CWC22 SOD1 SCAF11 NME1 ATM CDC20 SRSF11 TBL1XR1 RSL24D1 NDUFV1 SRSF7 BAG1 CCT6A ZEB1 ATP6V1F UQCRH CBFA2T3 FBL SRSF10 C12ORF10 VIRMA DUSP2 CCT3.\n"
     ]
    }
   ],
   "source": [
    "# Inspect a formatted sample\n",
    "print(\"--- Formatted Sample ---\")\n",
    "print(\"#----Model input:----#\")\n",
    "print(formatted_ds[0][\"model_input\"], \"\\n\")\n",
    "print(\"#----Response:----#\")\n",
    "print(formatted_ds[0][\"response\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c18db332",
   "metadata": {},
   "source": [
    "Now that our custom formatter is ready, we'll wrap our original `arrow_ds` in a `CSData` object. The `finetune` function will use this object and our custom formatter to prepare the data for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "28ab6da2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save directory for Huggingface dataset\n",
    "c2s_save_dir = \"/home/sr2464/scratch/C2S_API_Testing/Data/perturbation_tutorial/perturbation_tutorial\"\n",
    "c2s_save_name = \"jurkat_perturbation_c2s\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "fa776a5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "148d9b57e6d94845a51fa39c55e9846e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/21412 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CSData Object; Path=/home/sr2464/scratch/C2S_API_Testing/Data/perturbation_tutorial/perturbation_tutorial/jurkat_perturbation_c2s, Format=arrow\n"
     ]
    }
   ],
   "source": [
    "csdata = cs.CSData.csdata_from_arrow(\n",
    "    arrow_dataset=arrow_ds,  # Regular cell sentence dataset put here, finetune() function will repeat the formatting with the prompt formatter\n",
    "    vocabulary=vocabulary,\n",
    "    save_dir=c2s_save_dir,\n",
    "    save_name=c2s_save_name,\n",
    "    dataset_backend=\"arrow\"\n",
    ")\n",
    "print(csdata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2f93526",
   "metadata": {},
   "source": [
    "# Load a Pretrained Cell2Sentence Model\n",
    "\n",
    "We will start with a C2S model that has already been pretrained on a diverse set of single-cell datasets. This provides a strong foundation of biological knowledge. The `C2S-Scale-Pythia-1b-pt` and newer C2S-Scale models are good general-purpose models to start from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2acb7e97",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name_or_path = \"vandijklab/C2S-Scale-Pythia-1b-pt\"\n",
    "save_dir = \"/home/sr2464/scratch/C2S_API_Testing/Cache_Dir/perturbation_tutorial\"\n",
    "save_name = \"perturbation_pythia_1B\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "94d1353e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n",
      "CSModel Object; Path=/home/sr2464/scratch/C2S_API_Testing/Cache_Dir/perturbation_tutorial/perturbation_pythia_1B\n"
     ]
    }
   ],
   "source": [
    "csmodel = cs.CSModel(\n",
    "    model_name_or_path=model_name_or_path,\n",
    "    save_dir=save_dir,\n",
    "    save_name=save_name\n",
    ")\n",
    "print(csmodel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08006eaf",
   "metadata": {},
   "source": [
    "# Finetune for Perturbation Prediction\n",
    "\n",
    "Now, we'll finetune our model on the perturbation prediction task. We'll define our `TrainingArguments` and then call the `finetune()` method, passing in our `csdata` object and the `PerturbationPromptFormatter` instance we created.\n",
    "\n",
    "For this tutorial, we'll run for a small number of steps (`max_steps=500`). For a full finetuning run, you would typically train for several epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4a09042d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/sr2464/scratch/C2S_API_Testing/Cache_Dir/perturbation_tutorial/2025-10-14-23_33_53_finetune_perturbation_prediction\n"
     ]
    }
   ],
   "source": [
    "datetimestamp = datetime.now().strftime('%Y-%m-%d-%H_%M_%S')\n",
    "output_dir = os.path.join(csmodel.save_dir, f\"{datetimestamp}_finetune_{task_name}\")\n",
    "print(output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4d6af62c",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(output_dir):\n",
    "    os.mkdir(output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "2e28b043",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/gpfs/radev/home/sr2464/.conda/envs/cell2sentence2/lib/python3.8/site-packages/transformers/training_args.py:1559: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n"
     ]
    }
   ],
   "source": [
    "train_args = TrainingArguments(\n",
    "    bf16=True,\n",
    "    per_device_train_batch_size=2,\n",
    "    per_device_eval_batch_size=2,\n",
    "    gradient_accumulation_steps=4,\n",
    "    learning_rate=1e-5,\n",
    "    logging_steps=50,\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    num_train_epochs=1, \n",
    "    eval_steps=50,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    save_steps=100,\n",
    "    save_strategy=\"steps\",\n",
    "    output_dir=output_dir,\n",
    "    max_steps=500  # Shortened for tutorial purposes\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "04071118",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Grouping cells by perturbation and identifying global controls...\n",
      "Found 12013 control cells.\n",
      "Found 68 unique perturbations.\n",
      "Creating control-perturbed pairs...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 68/68 [00:03<00:00, 20.78it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reloading model from path on disk: /home/sr2464/scratch/C2S_API_Testing/Cache_Dir/perturbation_tutorial/perturbation_pythia_1B\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9318ab9f03604a26a3bce836398392fc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=3):   0%|          | 0/9399 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/gpfs/radev/home/sr2464/Desktop/cell2sentence/src/cell2sentence/csmodel.py:210: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `Trainer.__init__`. Use `processing_class` instead.\n",
      "Detected kernel version 4.18.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting training. Output directory: /home/sr2464/scratch/C2S_API_Testing/Cache_Dir/perturbation_tutorial/2025-10-14-23_33_53_finetune_perturbation_prediction\n",
      "Selecting 500 samples of eval dataset to shorten validation loop.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "max_steps is given, it will override any value given in num_train_epochs\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m The `run_name` is currently set to the same value as `TrainingArguments.output_dir`. If this was not intended, please specify a different run name by setting the `TrainingArguments.run_name` parameter.\n",
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33msyeda5688\u001b[0m (\u001b[33msyed-a-rizvi\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.22.2 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.16.6"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/sr2464/Desktop/cell2sentence/tutorials/wandb/run-20251014_233441-fa2rexvg</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/syed-a-rizvi/huggingface/runs/fa2rexvg' target=\"_blank\">/home/sr2464/scratch/C2S_API_Testing/Cache_Dir/perturbation_tutorial/2025-10-14-23_33_53_finetune_perturbation_prediction</a></strong> to <a href='https://wandb.ai/syed-a-rizvi/huggingface' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/syed-a-rizvi/huggingface' target=\"_blank\">https://wandb.ai/syed-a-rizvi/huggingface</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/syed-a-rizvi/huggingface/runs/fa2rexvg' target=\"_blank\">https://wandb.ai/syed-a-rizvi/huggingface/runs/fa2rexvg</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='13' max='500' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [ 13/500 00:07 < 05:49, 1.39 it/s, Epoch 0.01/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finetuning completed. Updated model saved to disk at: /home/sr2464/scratch/C2S_API_Testing/Cache_Dir/perturbation_tutorial/2025-10-14-23_33_53_finetune_perturbation_prediction\n"
     ]
    }
   ],
   "source": [
    "csmodel.fine_tune(\n",
    "    csdata=csdata,\n",
    "    task=task_name,\n",
    "    train_args=train_args,\n",
    "    loss_on_response_only=True, # We only want to calculate loss on the predicted sentence\n",
    "    top_k_genes=200,  # Use top 200 genes for this example, normally would use full cell sentence (all nonzero expressed genes) if possible\n",
    "    prompt_formatter=prompt_formatter  # Pass in our custom prompt formatter\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7de08f82",
   "metadata": {},
   "source": [
    "# Generate Predictions with the Finetuned Model\n",
    "\n",
    "After finetuning, let's load our new model and use it to predict the response to a perturbation for a cell from our test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "89474061",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/home/sr2464/scratch/C2S_API_Testing/Cache_Dir/perturbation_tutorial/2025-10-14-23_33_53_finetune_perturbation_prediction/checkpoint-500'"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "final_ckpt_path = os.path.join(output_dir, \"checkpoint-500\")\n",
    "final_ckpt_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "38d63e84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/home/sr2464/scratch/C2S_API_Testing/Cache_Dir/perturbation_tutorial'"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "save_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "ebff900c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n"
     ]
    }
   ],
   "source": [
    "# Load the finetuned model (it's automatically saved to csmodel.model_name_or_path)\n",
    "finetuned_model = cs.CSModel(\n",
    "    model_name_or_path=final_ckpt_path, # Path is updated after finetuning\n",
    "    save_dir=save_dir,\n",
    "    save_name=\"perturbation_predictor_finetuned_final\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "ca22f14f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/home/sr2464/scratch/C2S_API_Testing/Cache_Dir/perturbation_tutorial/perturbation_predictor_finetuned_final'"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "finetuned_model.save_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "cf85b83c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cuda'"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "7c36e4b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_model = AutoModelForCausalLM.from_pretrained(\n",
    "    finetuned_model.save_path,\n",
    "    cache_dir=os.path.join(csmodel.save_dir, \".cache\"),\n",
    "    trust_remote_code=True\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c828d229",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['train', 'val', 'test'])"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load dataset split done in finetune() function, saved to output directory\n",
    "with open(os.path.join(output_dir, 'data_split_indices_dict.pkl'), 'rb') as f:\n",
    "    data_split_indices_dict = pickle.load(f)\n",
    "\n",
    "data_split_indices_dict.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "54a9c647",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[7, 29, 33, 35, 51, 54, 65, 70, 115, 116]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Print a few indices of test samples\n",
    "data_split_indices_dict['test'][:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "6a54bcda",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['sample_type', 'model_input', 'response'],\n",
       "    num_rows: 10\n",
       "})"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Select a few unseen samples\n",
    "formatted_test_ds = formatted_ds.select(data_split_indices_dict['test'][:10])\n",
    "formatted_test_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "1feda129",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Inference Prompt ---\n",
      "Given the following cell sentence of 200 expressed genes representing a cell's basal state, predict the cell sentence after applying the perturbation: EIF4B.\n",
      "Control cell sentence: ACTB MT-ATP6 MT-CO3 MT-CO2 TMSB4X TUBA1B MT-ND4 HSP90AA1 HIST1H4C TMSB10 MT-CYB ACTG1 TUBB RACK1 PFN1 H3F3A ARHGDIB YBX1 MT-ND1 CFL1 H2AFZ B2M FTH1 HSP90AB1 CENPF UBA52 MIF MT-ND3 LDHB NCL PPIA SERF2 HNRNPA2B1 EEF1B2 CDK6 PSMA7 HSPD1 TOP2A KPNA2 GSTP1 BTF3 MT-ND2 CD3D SET CORO1A UBB ANP32B HNRNPU STMN1 EIF4A1 CCT3 EEF1G EEF2 SOX4 EIF1 BANF1 HNRNPD ADA COX4I1 MKI67 COX6C SUMO2 LCP1 SLC25A3 PKM CHCHD2 HINT1 OAZ1 NDUFS6 ENO1 TUBB4B HMGB2 MACF1 NME1 LSM4 SELENOH HSPE1 ERH NUCKS1 HMGN1 DDX5 NDUFB10 HIST1H1E GPX4 NDUFA13 GTF3A MT-ND6 PTGES3 ARL6IP1 RBMX PGK1 CTCF ATP5MC3 PSMB2 SLC25A5 PRDX5 COX7C SPTBN1 BPTF ANXA1 UBE2C PRDX1 SRSF10 NDUFA12 SFPQ EIF3G TOMM22 SIVA1 PHB2 CCT2 PA2G4 PCBP2 SUB1 CCT6A RHOA ASPM YWHAE ATP5MC1 XRCC5 SMC4 SRSF3 YBX3 HNRNPAB ATP5F1B CSDE1 SNRPB IFI16 SNRPD1 PAFAH1B3 MSN XRCC6 COX6A1 COX8A ATP5F1E ALYREF CBX3 MYL6 EIF3E RBM39 RAC1 UBE2D2 C1QBP KHDRBS1 CCT4 JPT1 RAD21 NME2 PSME2 SEPTIN7 SNU13 CKLF NMRAL1 NOP53 EIF3I EIF3D SAP18 SON HTATSF1 PSMA3 SEC61B AIP RBM8A C11ORF58 PRRC2C OSTC NDUFAB1 NDUFA4 SELENOW HNRNPA3 SRSF7 H3F3B TOMM20 SRSF9 HNRNPK CCT8 ARPC2 ATP5F1A ATP5F1C PHB EIF4B NUCB2 SSBP1 CD3G TMEM50A HNRNPM BUB3 SRSF2 PSMC5 KDM5A SNRNP25 PAK2 MYL12A NIFK SNRPF SLC25A6 KIF14 PRMT1 MTDH SRRM2 ATP5PF.\n",
      "\n",
      "Perturbed cell sentence:\n"
     ]
    }
   ],
   "source": [
    "# Select a sample from the test set for inference\n",
    "sample_idx = 0\n",
    "inference_prompt = formatted_test_ds[sample_idx]['model_input']\n",
    "ground_truth_response = formatted_test_ds[sample_idx]['response']\n",
    "\n",
    "print(\"--- Inference Prompt ---\")\n",
    "print(inference_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "96d8dd95",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the prediction\n",
    "predicted_response = finetuned_model.generate_from_prompt(\n",
    "    model=final_model,\n",
    "    prompt=inference_prompt,\n",
    "    max_num_tokens=800 # max number of tokens to generate, ~4 tokens per gene\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "c807b03c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--- Ground Truth Perturbed Cell ---\n",
      "MT-ATP6 MT-CO3 MT-CO2 MT-ND4 MT-CYB MT-ND1 MT-ND2 ACTB HSP90AA1 TMSB4X YBX1 EEF1B2 MT-ND3 RACK1 HSP90AB1 EEF1G MIF NME2 HIST1H4C TUBA1B NCL TUBB ADA ENO1 STMN1 H2AFZ PFN1 H3F3A CFL1 LDHB HINT1 HSPD1 C1QBP HSPE1 UBA52 SERF2 ACTG1 PPIA B2M CALR HNRNPA2B1 ARHGDIB GSTP1 SET MT-ND5 BTF3 CCT2 CHCHD2 NUCB2 XRCC5 PGK1 CD3D HNRNPU SUMO2 PPP1R14B HNRNPD UQCRH FDPS ALYREF SIVA1 DNAJA1 SLC25A6 ARPC2 FTL TYMS DUT COX4I1 SNRPB DDX5 PRKDC SLC25A3 PSMA7 CD3G SLC25A5 UBC ATP5F1E MYL6 ATP5F1B CCT6A PCLAF CDK6 H3F3B EIF4A1 EEF2 HMGB2 GUK1 THRAP3 HNRNPDL SERBP1 FABP5 EIF2S2 NUCKS1 HNRNPA3 HMGN1 COX7C SFPQ NDUFA13 HSPA8 OAZ1 MARCKSL1 DEK SELENOH FTH1 SRM SNRPF EIF1 DYNLL1 XRCC6 MT-ND6 RSL1D1 ANP32B EIF3E PSMA3 NDUFS5 ERH ATP5MC3 ARPC3 GLUL YBX3 FUS SNRPD1 YWHAB ATP5MC1 PRDX1 PSMA4 SOX4 NDUFAB1 PPA1 EIF5 NUDC STOML2 NME4 SRSF7 MZB1 H1FX NOP56 TPM4 NME1 NASP EIF3I ATP6V1F SF3B2 TRIR CCNI ARPC1B COX6C SLIRP SNRPC BRK1 ARPC5 ATP5MF SRSF3 CD7 HNRNPR GNAI2 CARHSP1 PPIG ATP5MG PTGES3 PFDN2 SSBP1 C9ORF16 PAICS GPX4 UBB C12ORF57 HIST1H1B ANP32A PRDX5 APRT PKM PFDN5 NHP2 SUB1 CORO1A LSM2 HNRNPC RRM2 PNN EPRS SNRPA1 HSP90B1 MCM3 CCT8 PRRC2C SKP1 RNASEH2B CIAO2B PRMT1 RAC1 SRRT HNRNPF VDAC3 ISG15 NAA10 RRP1B CCT3 MYL12B NAE1 EMC6.\n",
      "\n",
      "--- Predicted Perturbed Cell ---\n",
      "MT-ATP6 MT-CO3 MT-CO2 MT-CYB MT-ND4 TMSB4X ACTB MT-ND1 MT-ND3 MT-ND2 HSP90AA1 HSP90AB1 RACK1 YBX1 MIF TUBA1B H3F3A H2AFZ HSPD1 NCL EEF1B2 FTH1 HNRNPA2B1 EEF1G LDHB HINT1 CHCHD2 NME2 BTF3 UBA52 EIF4A1 ACTG1 STMN1 TUBB PFN1 CFL1 EIF1 PPIA GSTP1 PRDX1 HNRNPU SET B2M HSPE1 ARHGDIB ENO1 SERF2 CD3D MT-ND5 HNRNPD HNRNPC C1QBP FTL SLC25A3 COX4I1 SLC25A5 ATP5MC3 SRSF3 PPP1R14B SIVA1 HNRNPA3 CCT3 ANP32B SFPQ EIF3E HNRNPR PSMB1 ATP5MC2 CCT6A PSMB2 NUCB2 PSMA7 PSMB3 SNRPD1 PSMB6 ATP5F1B PPA1 UQCRH SNRPB NME1 SRSF7 PSMB7 NDUFS5 HNRNPK OAZ1 PFDN5 TMSB10 ADA GTF3A ATP5F1E SERBP1 COX7C NDUFA13 PSMB5 CCT2 ATP5F1D GUK1 EIF3A NDUFB10 FABP5 PSMB3 NDUFA4 ATP5F1A NUDC POMP HNRNPDL UBB PSMB6 NDUFS6 COX6A1 SRSF2 EIF3M SNRPE PSMA4 PSMD7 GADD45GIP1 ATP5MG PRMT1 NOP53 PSMD2 UQCRB YWHAE HNRNPM PSMD1 NDUFB11 EIF3H TOMM6 PSMD11 PSMC5 SRSF9 PSMC3 NDUFA11 NOP56 EIF3I EIF3G ATP5PO SNRPF RBM8A PSMD4 TOMM22 RTRAF SRSF11 PSMC1 NDUFAB1 HNRNPAB PSMD13 PSMA2 NUCKS1 UQCR10 RSL1D1 PSMB4 TOMM5 GYPC PFDN2 PSMA3 NDUFA12 NOP10 EIF4G2 SELENOH UBE2I RBM3 PSMD12 PSMC4 HNRNPD SRM PSMB8 EIF3L RBMX PSMD6 SLC25A6 TOMM20 UBE2L3 PSMB5 RBM25 EIF3D TMA7 PNN RBM39 UQCR11 GTF3C6 PPP1CA EIF2S2 PSMC2 NUDT8 PSMD11 RBM8B RBM3B TOMM40 RAC1 TUFM EIF4B PSMB7 RSL24D1 GTF2A2 NOL7 PSMD12.\n"
     ]
    }
   ],
   "source": [
    "print(\"\\n--- Ground Truth Perturbed Cell ---\")\n",
    "print(ground_truth_response)\n",
    "print(\"\\n--- Predicted Perturbed Cell ---\")\n",
    "print(predicted_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce20628b",
   "metadata": {},
   "source": [
    "# Conclusion\n",
    "\n",
    "In this tutorial, you learned how to finetune a Cell2Sentence model for a custom task: perturbation response prediction. By creating a `PerturbationPromptFormatter`, we were able to structure our data into control-perturbation-response triplets, enabling the model to learn the complex transcriptional changes that occur upon perturbation.\n",
    "\n",
    "This approach is highly flexible and can be adapted to various experimental designs. The finetuned model can now be used for in-silico experiments, such as virtual screening of genetic perturbations or predicting the effect of new compounds, significantly accelerating the pace of biological discovery."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f1296f4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cell2sentence2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
